{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Routing\n",
        "\n",
        "Routing classifies an input and directs it to a specialized followup task. This workflow allows for separation of concerns, and building more specialized prompts. Without this workflow, optimizing for one kind of input can hurt performance on other inputs."
      ],
      "metadata": {
        "id": "Zd3bcB169Wbk"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "![image.png]()"
      ],
      "metadata": {
        "id": "rc42GCfBAE0M"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "**When to use this workflow:** Routing works well for complex tasks where there are distinct categories that are better handled separately, and where classification can be handled accurately, either by an LLM or a more traditional classification model/algorithm."
      ],
      "metadata": {
        "id": "60CWIAWK9mw2"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Install Packages"
      ],
      "metadata": {
        "id": "aLLwj3mi9I17"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install -Uq openai-agents"
      ],
      "metadata": {
        "id": "lv55yrZP9B4f"
      },
      "execution_count": 10,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import nest_asyncio\n",
        "nest_asyncio.apply()"
      ],
      "metadata": {
        "id": "g8wpbPRd9DFl"
      },
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Config"
      ],
      "metadata": {
        "id": "ayJMbJgR9R5C"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from agents import (\n",
        "    AsyncOpenAI,\n",
        "    OpenAIChatCompletionsModel\n",
        ")\n",
        "from google.colab import userdata\n"
      ],
      "metadata": {
        "id": "gf_Lz2cz9Ebc"
      },
      "execution_count": 12,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "gemini_api_key = userdata.get(\"GEMINI_API_KEY\")\n",
        "\n",
        "\n",
        "# Check if the API key is present; if not, raise an error\n",
        "if not gemini_api_key:\n",
        "    raise ValueError(\"GEMINI_API_KEY is not set. Please ensure it is defined in your .env file.\")\n",
        "\n",
        "#Reference: https://ai.google.dev/gemini-api/docs/openai\n",
        "external_client = AsyncOpenAI(\n",
        "    api_key=gemini_api_key,\n",
        "    base_url=\"https://generativelanguage.googleapis.com/v1beta/openai/\",\n",
        ")\n",
        "\n",
        "model = OpenAIChatCompletionsModel(\n",
        "    model=\"gemini-2.0-flash\",\n",
        "    openai_client=external_client\n",
        ")"
      ],
      "metadata": {
        "id": "MQMt8cIv9F2V"
      },
      "execution_count": 13,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from agents import set_default_openai_client, set_tracing_disabled\n",
        "set_default_openai_client(external_client)\n",
        "set_tracing_disabled(True)"
      ],
      "metadata": {
        "id": "j5e69whK9HlT"
      },
      "execution_count": 14,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Implement Routing Pattern\n",
        "\n",
        "This shows the routing. The triage agent receives the first message, and\n",
        "then hands off to the appropriate agent based on the language of the request. Responses are\n",
        "streamed to the user.\n",
        "\n",
        "------------------------------------------------"
      ],
      "metadata": {
        "id": "fM1xyrpyAZh4"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 18,
      "metadata": {
        "id": "qfEZwzLU7skr"
      },
      "outputs": [],
      "source": [
        "import asyncio\n",
        "import uuid\n",
        "\n",
        "from openai.types.responses import ResponseContentPartDoneEvent, ResponseTextDeltaEvent\n",
        "\n",
        "from agents import Agent, RawResponsesStreamEvent, Runner, TResponseInputItem, trace\n",
        "\n",
        "french_agent = Agent(\n",
        "    name=\"french_agent\",\n",
        "    instructions=\"You only speak French\",\n",
        "    model=model\n",
        ")\n",
        "\n",
        "spanish_agent = Agent(\n",
        "    name=\"spanish_agent\",\n",
        "    instructions=\"You only speak Spanish\",\n",
        "    model=model\n",
        ")\n",
        "\n",
        "english_agent = Agent(\n",
        "    name=\"english_agent\",\n",
        "    instructions=\"You only speak English\",\n",
        "    model=model\n",
        ")\n",
        "\n",
        "triage_agent = Agent(\n",
        "    name=\"triage_agent\",\n",
        "    instructions=\"Handoff to the appropriate agent based on the language of the request.\",\n",
        "    handoffs=[french_agent, spanish_agent, english_agent],\n",
        "    model=model\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "async def main():\n",
        "    msg = input(\"Hi! We speak French, Spanish and English. How can I help? \")\n",
        "    agent = triage_agent\n",
        "    inputs: list[TResponseInputItem] = [{\"content\": msg, \"role\": \"user\"}]\n",
        "\n",
        "    while True:\n",
        "        # Each conversation turn is a single trace. Normally, each input from the user would be an\n",
        "        # API request to your app, and you can wrap the request in a trace()\n",
        "        result = Runner.run_streamed(\n",
        "            agent,\n",
        "            input=inputs,\n",
        "        )\n",
        "        async for event in result.stream_events():\n",
        "            if not isinstance(event, RawResponsesStreamEvent):\n",
        "                continue\n",
        "            data = event.data\n",
        "            if isinstance(data, ResponseTextDeltaEvent):\n",
        "                print(data.delta, end=\"\", flush=True)\n",
        "            elif isinstance(data, ResponseContentPartDoneEvent):\n",
        "                print(\"\\n\")\n",
        "\n",
        "        inputs = result.to_input_list()\n",
        "        print(\"\\n\")\n",
        "\n",
        "        user_msg = input(\"Enter a message: \")\n",
        "        if user_msg == \"exit\":\n",
        "          break\n",
        "        inputs.append({\"content\": user_msg, \"role\": \"user\"})\n",
        "        agent = result.current_agent"
      ],
      "metadata": {
        "id": "AAuCtjMLBCcm"
      },
      "execution_count": 23,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "asyncio.run(main())"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Y42jDAEAA2kJ",
        "outputId": "8e37e873-2a2c-4b89-c7a6-0b8f2960416a"
      },
      "execution_count": 24,
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Hi! We speak French, Spanish and English. How can I help? I want to book a flight\n",
            "Okay, I can help you with that!  To find the best flight for you, I need some information.  Please tell me:\n",
            "\n",
            "*   **Where are you flying from?** (Departure city/airport)\n",
            "*   **Where are you flying to?** (Destination city/airport)\n",
            "*   **What date would you like to depart?**\n",
            "*   **What date would you like to return?** (If it's a round trip)\n",
            "*   **How many passengers?** (Adults, children, infants)\n",
            "*   **Do you have any preferred airlines?**\n",
            "*   **Do you have any preferred class of service?** (Economy, Business, First Class)\n",
            "*   **Do you have a maximum budget?** (Optional, but helpful)\n",
            "\n",
            "\n",
            "\n",
            "\n",
            "Enter a message: Hmmm\n",
            "Is there something I can clarify or help you with? Are you unsure about any of the information I requested? Just let me know!\n",
            "\n",
            "\n",
            "\n",
            "\n",
            "Enter a message: exit\n"
          ]
        }
      ]
    }
  ]
}